import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from PIL import Image
from torchtoolbox.transform import Cutout
import numpy as np
from timm.data.mixup import Mixup


batch_size = 400
num_workers = 2
learning_rate=0.001
epochs = 50  # 50轮
# 判断是否有GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 图像预处理 （数是官网找的）
mean = [0.50505644, 0.49160016, 0.4444035]  # [0.5070751592371323, 0.48654887331495095, 0.4409178433670343]
std = [0.25895253, 0.25304303, 0.27547014]  # [0.2673342858792401, 0.2564384629170883, 0.27615047132568404]

# 数据增广
transform_train = transforms.Compose([
    # transforms.Resize((64, 64)), # 从32*32到224*224
    transforms.RandomCrop(32, padding=4), # 随机裁剪
    transforms.RandomHorizontalFlip(),# 水平翻转
    transforms.RandomVerticalFlip(), # 垂直翻转
    transforms.RandomRotation(15), # 随机旋转
    transforms.ToTensor(),
    transforms.Normalize(mean, std),
    #Cutout(0.5), # 参数是遮挡的概率
    #transforms.RandomErasing(),# 随机擦除
])
transform_test = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize(mean, std)])


# CIFAR-100 数据集下载
train_dataset = torchvision.datasets.CIFAR100(root='data/',
                                             train=True,
                                             transform=transform_train,
                                             download=False)

test_dataset = torchvision.datasets.CIFAR100(root='data/',
                                            train=False,
                                            transform=transform_test)

# 数据载入
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           num_workers=num_workers,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          num_workers=num_workers,
                                          shuffle=False)




def gen_mean_std(dataset):
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=2)
    train = next(iter(dataloader))[0]
    mean = np.mean(train.numpy(), axis=(0, 2, 3))
    std = np.std(train.numpy(), axis=(0, 2, 3))
    return mean, std



if __name__=='__main__':

    cifar100 = torchvision.datasets.CIFAR100(root='./data', train=True, download=False, transform=transforms.Compose([transforms.ToTensor()]))
    mean, std = gen_mean_std(cifar100)
    print(mean, std)


